{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "260d6e52",
   "metadata": {},
   "source": [
    "# 测试回调函数"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5bd0f71",
   "metadata": {},
   "source": [
    "测试 TVM Relax 虚拟机（VM）中 Python 回调函数与 Relax 函数之间的交互机制，具体包括： \n",
    "1. 双向数据传递：测试 Relax 函数向Python回调函数传递张量，以及Python回调函数向Relax函数返回张量的能力。\n",
    "2. 异常传播：测试Python回调函数中抛出的异常是否能够正确传播到调用方，并且保留完整的堆栈跟踪信息。\n",
    "3. 多种执行模式：所有测试都在两种执行模式下运行 - bytecode（字节码）模式和compiled（编译）模式，以验证VM在不同执行策略下的行为一致性。   \n",
    "\n",
    "文件使用TVM的测试框架，通过参数化测试方式自动在不同配置下运行测试用例，确保Relax VM的回调机制能够正常工作。这些测试对于保证TVM Relax与Python之间的交互可靠性非常重要，特别是在需要混合执行Python代码和编译后代码的场景中。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "55c2e9d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "import tvm.testing\n",
    "\n",
    "from tvm.script import relax as R\n",
    "import numpy as np\n",
    "\n",
    "# 定义执行模式参数化列表，测试将在字节码模式和编译模式下分别运行\n",
    "EXEC_MODE = [\"bytecode\", \"compiled\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4191f53d",
   "metadata": {},
   "source": [
    "## 测试将张量从 Relax 函数传递给 Python 回调函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "20b228d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义Relax函数，该函数接受一个张量和一个回调函数作为参数\n",
    "@R.function\n",
    "def relax_func(\n",
    "    A: R.Tensor([16], \"int32\"),\n",
    "    callback: R.Callable([R.Tensor([16], \"int32\")], R.Tuple([])),\n",
    "):\n",
    "    # 将输入张量A乘以2\n",
    "    B = R.multiply(A, R.const(2))\n",
    "    # 调用回调函数，并传入计算结果B\n",
    "    _ = callback(B)\n",
    "    # 返回空元组\n",
    "    return R.tuple()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d2479264",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(exec_mode, target=\"llvm\", dev=tvm.cpu()):\n",
    "    # 构建Relax模块，根据指定的执行模式（bytecode或compiled）\n",
    "    ex = tvm.relax.build(\n",
    "        tvm.IRModule.from_expr(relax_func),\n",
    "        target=target,\n",
    "        exec_mode=exec_mode,\n",
    "    )\n",
    "    # 创建虚拟机实例来执行编译后的模块\n",
    "    vm = tvm.relax.VirtualMachine(ex, dev)\n",
    "\n",
    "    # 全局变量，用于存储回调函数接收到的张量\n",
    "    from_callback = None\n",
    "\n",
    "    # 定义自定义的Python回调函数\n",
    "    def custom_callback(arr):\n",
    "        # 非局部变量声明，允许在内部函数中修改外部函数的变量\n",
    "        nonlocal from_callback\n",
    "        # 保存接收到的张量\n",
    "        from_callback = arr\n",
    "\n",
    "    # 创建测试数据\n",
    "    np_A = np.arange(16, dtype=\"int32\")\n",
    "    # 将numpy数组转换为TVM NDArray\n",
    "    tvm_A = tvm.nd.array(np_A)\n",
    "\n",
    "    # 在虚拟机上执行Relax函数，传入张量和回调函数\n",
    "    vm[\"relax_func\"](tvm_A, custom_callback)\n",
    "\n",
    "    # 验证回调函数是否被正确调用，并且接收到了正确的结果\n",
    "    assert from_callback is not None\n",
    "    np.testing.assert_array_equal(np_A * 2, from_callback.numpy())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b0d2d66",
   "metadata": {},
   "outputs": [],
   "source": [
    "for exec_mode in EXEC_MODE:\n",
    "    test(exec_mode, target=\"llvm\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "237d02b1",
   "metadata": {},
   "source": [
    "## 测试从Python回调函数生成张量并传递给Relax函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "892ef5de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义Relax函数，该函数接受一个无参数但返回张量的回调函数\n",
    "@R.function\n",
    "def relax_func(\n",
    "    callback: R.Callable([], R.Tensor([16], \"int32\")),\n",
    "):\n",
    "    # 调用回调函数获取张量\n",
    "    A = callback()\n",
    "    # 将获取的张量乘以2\n",
    "    B = R.multiply(A, R.const(2))\n",
    "    # 返回计算结果\n",
    "    return B\n",
    "target = \"llvm\"\n",
    "dev = tvm.cpu(0)\n",
    "for exec_mode in EXEC_MODE:\n",
    "    # 构建Relax模块\n",
    "    ex = tvm.relax.build(\n",
    "        tvm.IRModule.from_expr(relax_func),\n",
    "        target=target,\n",
    "        exec_mode=exec_mode,\n",
    "    )\n",
    "    # 创建虚拟机实例\n",
    "    vm = tvm.relax.VirtualMachine(ex, dev)\n",
    "\n",
    "    # 创建测试数据\n",
    "    np_A = np.arange(16, dtype=\"int32\")\n",
    "\n",
    "    # 定义自定义回调函数，返回测试数据的TVM NDArray\n",
    "    def custom_callback():\n",
    "        return tvm.nd.array(np_A)\n",
    "\n",
    "    # 在虚拟机上执行Relax函数，传入回调函数，并获取输出结果\n",
    "    output = vm[\"relax_func\"](custom_callback)\n",
    "\n",
    "    # 验证输出结果是否正确\n",
    "    np.testing.assert_array_equal(np_A * 2, output.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2c90b36",
   "metadata": {},
   "source": [
    "## 测试Python回调函数中的异常是否能正确传播，并保留完整的堆栈跟踪"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22c99fb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义Relax函数，接受一个返回张量的回调函数\n",
    "@R.function\n",
    "def relax_func(\n",
    "    callback: R.Callable([], R.Tensor([16], \"int32\")),\n",
    "):\n",
    "    # 调用回调函数获取张量\n",
    "    A = callback()\n",
    "    # 返回获取的张量\n",
    "    return A\n",
    "target = \"llvm\"\n",
    "dev = tvm.cpu(0)\n",
    "for exec_mode in EXEC_MODE:\n",
    "    # 构建Relax模块\n",
    "    ex = tvm.relax.build(\n",
    "        tvm.IRModule.from_expr(relax_func),\n",
    "        target=target,\n",
    "        exec_mode=exec_mode,\n",
    "    )\n",
    "    # 创建虚拟机实例\n",
    "    vm = tvm.relax.VirtualMachine(ex, dev)\n",
    "\n",
    "    # 定义一个会抛出异常的自定义回调函数\n",
    "    def custom_callback():\n",
    "        # 定义一个局部变量，用于测试堆栈中的局部变量信息\n",
    "        local_var = 42\n",
    "        # 抛出运行时错误\n",
    "        raise RuntimeError(\"Error thrown from callback\")\n",
    "\n",
    "    try:\n",
    "        # 在虚拟机上执行Relax函数，传入会抛出异常的回调函数\n",
    "        vm[\"relax_func\"](custom_callback)\n",
    "    except RuntimeError as err:\n",
    "        # 获取异常的堆栈跟踪\n",
    "        stack = err.__traceback__\n",
    "        # 遍历堆栈，找到最内层的栈帧（回调函数所在的栈帧）\n",
    "        while stack.tb_next is not None:\n",
    "            stack = stack.tb_next\n",
    "        frame = stack.tb_frame\n",
    "        # 验证最内层的栈帧是否来自Python回调函数\n",
    "        assert (\n",
    "            frame.f_code.co_filename.find(\"test_vm_callback_function.py\") != -1\n",
    "        ), \"Inner-most stack frame should be from Python callback\"\n",
    "\n",
    "    else:\n",
    "        # 如果没有捕获到异常，抛出错误\n",
    "        raise RuntimeError(\"Exception thrown in callback was not propagated to calling scope\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6056472",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
